import numpy as np
# import quimb as qu
import jax.numpy as jnp

from quimb.tensor.tensor_core import (
	Tensor, TensorNetwork, rand_uuid
)

G0 = jnp.diag(jnp.array([1.0, 0, 0, 0]))
G1 = jnp.diag(jnp.array([0.0, 1, 0, 0]))
G2 = jnp.diag(jnp.array([0.0, 0, 1, 0]))
G3 = jnp.diag(jnp.array([0.0, 0, 0, 1]))

G_o3 = jnp.array([G0, G1, G2, G3])

def splm_1qu_ptm(lam1s):
	lam1, lam2, lam3 = lam1s
	return jnp.diag(jnp.array([1.0, jnp.exp(-2.0 * (lam2 + lam3)), jnp.exp(-2.0 * (lam1 + lam3)), jnp.exp(-2.0 * (lam1 + lam2))]))

def splm_2qu_ptm(lam2s):
	lam11, lam12, lam13, lam21, lam22, lam23, lam31, lam32, lam33 = lam2s
	zeta00 = 1.0
	zeta01 = jnp.exp(-2.0 * (lam12 + lam13 + lam22 + lam23 + lam32 + lam33))
	zeta02 = jnp.exp(-2.0 * (lam11 + lam13 + lam21 + lam23 + lam31 + lam33))
	zeta03 = jnp.exp(-2.0 * (lam11 + lam12 + lam21 + lam22 + lam31 + lam32))
	zeta10 = jnp.exp(-2.0 * (lam21 + lam22 + lam23 + lam31 + lam32 + lam33))
	zeta11 = jnp.exp(-2.0 * (lam12 + lam13 + lam21 + lam31))
	zeta12 = jnp.exp(-2.0 * (lam11 + lam13 + lam22 + lam32))
	zeta13 = jnp.exp(-2.0 * (lam11 + lam12 + lam23 + lam33))
	zeta20 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam31 + lam32 + lam33))
	zeta21 = jnp.exp(-2.0 * (lam11 + lam22 + lam23 + lam31))
	zeta22 = jnp.exp(-2.0 * (lam12 + lam21 + lam23 + lam32))
	zeta23 = jnp.exp(-2.0 * (lam13 + lam21 + lam22 + lam33))
	zeta30 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam21 + lam22 + lam23))
	zeta31 = jnp.exp(-2.0 * (lam11 + lam21 + lam32 + lam33))
	zeta32 = jnp.exp(-2.0 * (lam12 + lam22 + lam31 + lam33))
	zeta33 = jnp.exp(-2.0 * (lam13 + lam23 + lam31 + lam32))
	F0 = jnp.diag(jnp.array([zeta00, zeta10, zeta20, zeta30]))
	F1 = jnp.diag(jnp.array([zeta01, zeta11, zeta21, zeta31]))
	F2 = jnp.diag(jnp.array([zeta02, zeta12, zeta22, zeta32]))
	F3 = jnp.diag(jnp.array([zeta03, zeta13, zeta23, zeta33]))
	return jnp.array([F0, F1, F2, F3])

def SPLM_ptm_mpo(
		N_qubit: int,
		lam1_mat, lam2_mat,
		upper_ind_id="N{}", lower_ind_id="N'{}",
		site_tag_id="I{}",
		**tn_opts
):
	assert len(lam1_mat) == len(lam2_mat) + 1 == N_qubit, f"mismatch length {N_qubit}"

	def gen_tensors():
		upper_inds, lower_inds = map(upper_ind_id.format, range(N_qubit)), map(lower_ind_id.format, range(N_qubit))
		site_tags = tuple(map(site_tag_id.format, range(N_qubit)))

		k1_ix, n_ix, site_tag = rand_uuid(), rand_uuid(), site_tags[0]
		t1 = Tensor(splm_1qu_ptm(lam1_mat[0]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
		t2 = Tensor(splm_2qu_ptm(lam2_mat[0]), inds=(n_ix, next(upper_inds), k1_ix), tags=[site_tag])
		# yield t1 @ t2
		yield t1
		yield t2
		p_ix = n_ix
		for i in range(1, N_qubit - 1):
			k1_ix, k2_ix, n_ix, site_tag = rand_uuid(), rand_uuid(), rand_uuid(), site_tags[i]
			t1 = Tensor(splm_1qu_ptm(lam1_mat[i]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
			if i % 2:
				t2 = Tensor(G_o3, inds=(p_ix, k2_ix, k1_ix), tags=[site_tag])
				t3 = Tensor(splm_2qu_ptm(lam2_mat[i]), inds=(n_ix, next(upper_inds), k2_ix), tags=[site_tag])
			else:
				t2 = Tensor(splm_2qu_ptm(lam2_mat[i]), inds=(n_ix, k2_ix, k1_ix), tags=[site_tag])
				t3 = Tensor(G_o3, inds=(p_ix, next(upper_inds), k2_ix), tags=[site_tag])
			# yield t1 @ t2 @ t3
			yield t1
			yield t2
			yield t3
			p_ix = n_ix
		k1_ix, site_tag = rand_uuid(), site_tags[-1]
		t1 = Tensor(splm_1qu_ptm(lam1_mat[-1]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
		t2 = Tensor(G_o3, inds=(p_ix, next(upper_inds), k1_ix), tags=[site_tag])
		# yield t1 @ t2
		yield t1
		yield t2

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def inv_splm_1qu_ptm(lam1s):
	lam1, lam2, lam3 = lam1s
	return np.diag([1.0, np.exp(2.0 * (lam2 + lam3)), np.exp(2.0 * (lam1 + lam3)), np.exp(2.0 * (lam1 + lam2))])

def inv_splm_2qu_ptm(lam2s):
	lam11, lam12, lam13, lam21, lam22, lam23, lam31, lam32, lam33 = lam2s
	zeta00 = 1.0
	zeta01 = np.exp(2.0 * (lam12 + lam13 + lam22 + lam23 + lam32 + lam33))
	zeta02 = np.exp(2.0 * (lam11 + lam13 + lam21 + lam23 + lam31 + lam33))
	zeta03 = np.exp(2.0 * (lam11 + lam12 + lam21 + lam22 + lam31 + lam32))
	zeta10 = np.exp(2.0 * (lam21 + lam22 + lam23 + lam31 + lam32 + lam33))
	zeta11 = np.exp(2.0 * (lam12 + lam13 + lam21 + lam31))
	zeta12 = np.exp(2.0 * (lam11 + lam13 + lam22 + lam32))
	zeta13 = np.exp(2.0 * (lam11 + lam12 + lam23 + lam33))
	zeta20 = np.exp(2.0 * (lam11 + lam12 + lam13 + lam31 + lam32 + lam33))
	zeta21 = np.exp(2.0 * (lam11 + lam22 + lam23 + lam31))
	zeta22 = np.exp(2.0 * (lam12 + lam21 + lam23 + lam32))
	zeta23 = np.exp(2.0 * (lam13 + lam21 + lam22 + lam33))
	zeta30 = np.exp(2.0 * (lam11 + lam12 + lam13 + lam21 + lam22 + lam23))
	zeta31 = np.exp(2.0 * (lam11 + lam21 + lam32 + lam33))
	zeta32 = np.exp(2.0 * (lam12 + lam22 + lam31 + lam33))
	zeta33 = np.exp(2.0 * (lam13 + lam23 + lam31 + lam32))
	F0 = np.diag([zeta00, zeta10, zeta20, zeta30])
	F1 = np.diag([zeta01, zeta11, zeta21, zeta31])
	F2 = np.diag([zeta02, zeta12, zeta22, zeta32])
	F3 = np.diag([zeta03, zeta13, zeta23, zeta33])
	return np.array([F0, F1, F2, F3])

def inv_SPLM_ptm_mpo(
		N_qubit: int,
		lam1_mat, lam2_mat,
		upper_ind_id="N{}", lower_ind_id="N'{}",
		site_tag_id="I{}",
		**tn_opts
):
	assert len(lam1_mat) == len(lam2_mat) + 1 == N_qubit, f"mismatch length {N_qubit}"

	def gen_tensors():
		upper_inds, lower_inds = map(upper_ind_id.format, range(N_qubit)), map(lower_ind_id.format, range(N_qubit))
		site_tags = tuple(map(site_tag_id.format, range(N_qubit)))

		k1_ix, n_ix, site_tag = rand_uuid(), rand_uuid(), site_tags[0]
		t1 = Tensor(inv_splm_1qu_ptm(lam1_mat[0]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
		t2 = Tensor(inv_splm_2qu_ptm(lam2_mat[0]), inds=(n_ix, next(upper_inds), k1_ix), tags=[site_tag])
		yield t1 @ t2
		p_ix = n_ix
		for i in range(1, N_qubit - 1):
			k1_ix, k2_ix, n_ix, site_tag = rand_uuid(), rand_uuid(), rand_uuid(), site_tags[i]
			t1 = Tensor(inv_splm_1qu_ptm(lam1_mat[i]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
			if i % 2:
				t2 = Tensor(G_o3, inds=(p_ix, k2_ix, k1_ix), tags=[site_tag])
				t3 = Tensor(inv_splm_2qu_ptm(lam2_mat[i]), inds=(n_ix, next(upper_inds), k2_ix), tags=[site_tag])
			else:
				t2 = Tensor(inv_splm_2qu_ptm(lam2_mat[i]), inds=(n_ix, k2_ix, k1_ix), tags=[site_tag])
				t3 = Tensor(G_o3, inds=(p_ix, next(upper_inds), k2_ix), tags=[site_tag])
			yield t1 @ t2 @ t3
			p_ix = n_ix
		k1_ix, site_tag = rand_uuid(), site_tags[-1]
		t1 = Tensor(inv_splm_1qu_ptm(lam1_mat[-1]), inds=(k1_ix, next(lower_inds)), tags=[site_tag])
		t2 = Tensor(G_o3, inds=(p_ix, next(upper_inds), k1_ix), tags=[site_tag])
		yield t1 @ t2

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

# from quimb.tensor import PTensor
# def SPLM_ptm_mpo_p(
# 		N_qubit: int,
# 		lam1_mat, lam2_mat,
# 		upper_ind_id="N{}", lower_ind_id="N'{}",
# 		site_tag_id="I{}",
# 		site_p_tag='Nv', site_c_tag='Nc',
# 		**tn_opts
# ):
# 	def splm_1qu_ptm_p(lam1s):
# 		lam1, lam2, lam3 = lam1s
# 		return jnp.diag(
# 				[1.0, jnp.exp(-2.0 * (lam2 + lam3)), jnp.exp(-2.0 * (lam1 + lam3)), jnp.exp(-2.0 * (lam1 + lam2))])
#
# 	def splm_2qu_ptm_p(lam2s):
# 		lam11, lam12, lam13, lam21, lam22, lam23, lam31, lam32, lam33 = lam2s
# 		zeta00 = 1.0
# 		zeta01 = jnp.exp(-2.0 * (lam12 + lam13 + lam22 + lam23 + lam32 + lam33))
# 		zeta02 = jnp.exp(-2.0 * (lam11 + lam13 + lam21 + lam23 + lam31 + lam33))
# 		zeta03 = jnp.exp(-2.0 * (lam11 + lam12 + lam21 + lam22 + lam31 + lam32))
# 		zeta10 = jnp.exp(-2.0 * (lam21 + lam22 + lam23 + lam31 + lam32 + lam33))
# 		zeta11 = jnp.exp(-2.0 * (lam12 + lam13 + lam21 + lam31))
# 		zeta12 = jnp.exp(-2.0 * (lam11 + lam13 + lam22 + lam32))
# 		zeta13 = jnp.exp(-2.0 * (lam11 + lam12 + lam23 + lam33))
# 		zeta20 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam31 + lam32 + lam33))
# 		zeta21 = jnp.exp(-2.0 * (lam11 + lam22 + lam23 + lam31))
# 		zeta22 = jnp.exp(-2.0 * (lam12 + lam21 + lam23 + lam32))
# 		zeta23 = jnp.exp(-2.0 * (lam13 + lam21 + lam22 + lam33))
# 		zeta30 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam21 + lam22 + lam23))
# 		zeta31 = jnp.exp(-2.0 * (lam11 + lam21 + lam32 + lam33))
# 		zeta32 = jnp.exp(-2.0 * (lam12 + lam22 + lam31 + lam33))
# 		zeta33 = jnp.exp(-2.0 * (lam13 + lam23 + lam31 + lam32))
# 		F0 = jnp.diag([zeta00, zeta10, zeta20, zeta30])
# 		F1 = jnp.diag([zeta01, zeta11, zeta21, zeta31])
# 		F2 = jnp.diag([zeta02, zeta12, zeta22, zeta32])
# 		F3 = jnp.diag([zeta03, zeta13, zeta23, zeta33])
# 		return jnp.array([F0, F1, F2, F3])
#
# 	def gen_tensors():
# 		upper_inds, lower_inds, site_tags = map(upper_ind_id.format, range(N_qubit)), \
# 			map(lower_ind_id.format, range(N_qubit)), map(site_tag_id.format, range(N_qubit))
#
# 		k1_ix, n_ix = rand_uuid(), rand_uuid()
# 		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
# 		t1 = PTensor(splm_1qu_ptm_p, lam1_mat[0], inds=(k1_ix, lower_ind), tags=[site_tag, site_p_tag])
# 		yield t1
# 		t2 = PTensor(splm_2qu_ptm_p, lam2_mat[0], inds=(n_ix, upper_ind, k1_ix), tags=[site_tag, site_p_tag])
# 		yield t2
# 		p_ix = n_ix
# 		for i in range(1, N_qubit - 1):
# 			k1_ix, k2_ix, n_ix = rand_uuid(), rand_uuid(), rand_uuid()
# 			upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
# 			t1 = PTensor(splm_1qu_ptm_p, lam1_mat[i], inds=(k1_ix, lower_ind), tags=[site_tag, site_p_tag])
# 			if i % 2:
# 				t2 = Tensor(G_o3, inds=(p_ix, k2_ix, k1_ix), tags=[site_tag, site_c_tag])
# 				t3 = PTensor(splm_2qu_ptm_p, lam2_mat[i], inds=(n_ix, upper_ind, k2_ix), tags=[site_tag, site_p_tag])
# 			else:
# 				t2 = PTensor(splm_2qu_ptm_p, lam2_mat[i], inds=(n_ix, k2_ix, k1_ix), tags=[site_tag, site_p_tag])
# 				t3 = Tensor(G_o3, inds=(p_ix, upper_ind, k2_ix), tags=[site_tag, site_c_tag])
# 			yield t1
# 			yield t2
# 			yield t3
# 			p_ix = n_ix
# 		k1_ix = rand_uuid()
# 		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
# 		t1 = PTensor(splm_1qu_ptm_p, lam1_mat[-1], inds=(k1_ix, lower_ind), tags=[site_tag, site_p_tag])
# 		t2 = Tensor(G_o3, inds=(p_ix, upper_ind, k1_ix), tags=[site_tag, site_c_tag])
# 		yield t1
# 		yield t2
#
# 	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)